package mn.idax.exchange.otc.util;

import mn.idax.exchange.auth.util.IpUtils;
import mn.idax.exchange.otc.annotation.RequestLimit;
import mn.idax.exchange.otc.enums.RequestLimitType;
import org.apache.commons.lang3.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import javax.servlet.http.HttpServletRequest;

/**
 * @Auther: dison
 * @Date: 2018/8/14 14:06
 * @Description:
 */
public class RequestLimitUtils {

    /**
     * 限流 脚本
     *
     * @return lua脚本
     */
    public static String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用不超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }

    public static String getRquestLimitKey(RequestLimit limitAnnotation,String className){
        String key;
        RequestLimitType limitType = limitAnnotation.limitType();
        switch (limitType) {
            case IP:
                HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
                key = "IP:" + IpUtils.getIpAddr(request).replaceAll(":",".");
                break;
            case CUSTOMER:
                key = "CUSTOMER:" +limitAnnotation.key();
                break;
            default:
                key = "CLASSNAME:" +StringUtils.upperCase(className);
        }
        return key;
    }
}
